import os
import json
import hashlib
from flask import Flask, request, jsonify, send_file
from werkzeug.utils import secure_filename
from flask_cors import CORS

app = Flask(__name__)
CORS(app)  # 启用CORS，允许跨域请求

# 配置
UPLOAD_FOLDER = 'uploads'
CHUNK_SIZE = 5 * 1024 * 1024  # 5MB，与前端保持一致
os.makedirs(UPLOAD_FOLDER, exist_ok=True)

# 检查文件是否已上传
@app.route('/check-file', methods=['POST'])
def check_file():
    data = request.json
    file_hash = data.get('hash')
    
    # 检查文件是否已存在
    file_path = os.path.join(UPLOAD_FOLDER, file_hash)
    if os.path.exists(file_path):
        return jsonify({
            'exists': True,
            'message': '文件已存在'
        })
    
    # 检查已上传的分片
    chunk_dir = os.path.join(UPLOAD_FOLDER, f'{file_hash}_chunks')
    uploaded_chunks = []
    if os.path.exists(chunk_dir):
        for filename in os.listdir(chunk_dir):
            if filename.startswith('chunk_'):
                chunk_index = int(filename.split('_')[1])
                uploaded_chunks.append(chunk_index)
    
    return jsonify({
        'exists': False,
        'uploaded_chunks': uploaded_chunks,
        'message': '准备上传'
    })

# 上传分片
@app.route('/upload-chunk', methods=['POST'])
def upload_chunk():
    file = request.files['file']
    chunk_index = int(request.form.get('chunk_index'))
    total_chunks = int(request.form.get('total_chunks'))
    file_hash = request.form.get('file_hash')
    original_filename = secure_filename(request.form.get('filename'))
    
    # 创建分片目录
    chunk_dir = os.path.join(UPLOAD_FOLDER, f'{file_hash}_chunks')
    os.makedirs(chunk_dir, exist_ok=True)
    
    # 保存分片
    chunk_path = os.path.join(chunk_dir, f'chunk_{chunk_index}')
    file.save(chunk_path)
    
    # 检查是否所有分片都已上传
    uploaded_chunks = []
    for filename in os.listdir(chunk_dir):
        if filename.startswith('chunk_'):
            uploaded_chunks.append(int(filename.split('_')[1]))
    
    # 如果所有分片都已上传，合并它们
    if len(uploaded_chunks) == total_chunks and sorted(uploaded_chunks) == list(range(total_chunks)):
        merge_chunks(file_hash, total_chunks, original_filename)
    
    return jsonify({
        'success': True,
        'chunk_index': chunk_index,
        'message': '分片上传成功'
    })

# 合并分片
def merge_chunks(file_hash, total_chunks, original_filename):
    chunk_dir = os.path.join(UPLOAD_FOLDER, f'{file_hash}_chunks')
    output_path = os.path.join(UPLOAD_FOLDER, file_hash)
    
    # 合并分片
    with open(output_path, 'wb') as outfile:
        for i in range(total_chunks):
            chunk_path = os.path.join(chunk_dir, f'chunk_{i}')
            with open(chunk_path, 'rb') as infile:
                outfile.write(infile.read())
            os.remove(chunk_path)  # 删除已合并的分片
    
    # 删除分片目录
    os.rmdir(chunk_dir)
    
    # 保存文件元数据
    metadata = {
        'hash': file_hash,
        'original_filename': original_filename,
        'size': os.path.getsize(output_path),
        'uploaded_at': str(datetime.now())
    }
    metadata_path = os.path.join(UPLOAD_FOLDER, f'{file_hash}.json')
    with open(metadata_path, 'w') as f:
        json.dump(metadata, f)
    
    return True

# 标记文件为已上传（用于秒传）
@app.route('/mark-uploaded', methods=['POST'])
def mark_uploaded():
    data = request.json
    file_hash = data.get('hash')
    original_filename = secure_filename(data.get('filename'))
    
    # 创建一个空文件表示已上传
    file_path = os.path.join(UPLOAD_FOLDER, file_hash)
    with open(file_path, 'w') as f:
        pass
    
    # 保存文件元数据
    metadata = {
        'hash': file_hash,
        'original_filename': original_filename,
        'size': 0,  # 秒传文件大小为0
        'uploaded_at': str(datetime.now()),
        'fast_upload': True
    }
    metadata_path = os.path.join(UPLOAD_FOLDER, f'{file_hash}.json')
    with open(metadata_path, 'w') as f:
        json.dump(metadata, f)
    
    return jsonify({
        'success': True,
        'message': '文件已标记为上传'
    })

# 获取文件列表
@app.route('/files', methods=['GET'])
def get_files():
    files = []
    for filename in os.listdir(UPLOAD_FOLDER):
        if not filename.endswith('.json') and not filename.endswith('_chunks'):
            metadata_path = os.path.join(UPLOAD_FOLDER, f'{filename}.json')
            if os.path.exists(metadata_path):
                with open(metadata_path, 'r') as f:
                    metadata = json.load(f)
                    files.append(metadata)
    
    return jsonify({
        'files': files,
        'count': len(files)
    })

# 下载文件
@app.route('/download/<file_hash>', methods=['GET'])
def download_file(file_hash):
    file_path = os.path.join(UPLOAD_FOLDER, file_hash)
    if not os.path.exists(file_path):
        return jsonify({'error': '文件不存在'}), 404
    
    # 获取原始文件名
    metadata_path = os.path.join(UPLOAD_FOLDER, f'{file_hash}.json')
    original_filename = file_hash
    if os.path.exists(metadata_path):
        with open(metadata_path, 'r') as f:
            metadata = json.load(f)
            original_filename = metadata.get('original_filename', file_hash)
    
    return send_file(file_path, as_attachment=True, download_name=original_filename)

if __name__ == '__main__':
    app.run(debug=True)    